import pandas as pd
import os
import sqlite3

class DataHelper:
    def __init__(self,filename='data.csv'):
        self.filename = filename
        self.dataMode=2 if filename.find('data.db')>=0 else 1

    def load_data(self):
        if self.dataMode == 1 :
            os.path.exists(self.filename)
            return pd.read_csv(self.filename) 
        if self.dataMode == 2:
            conn = sqlite3.connect(self.filename)  
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM users")    
            # 获取列名
            column_names = [description[0] for description in cursor.description]
            # 获取所有数据
            data = cursor.fetchall()
            # 将数据转换为 DataFrame
            df = pd.DataFrame(data, columns=column_names)
            # 打印 DataFrame
            print(df)
            return df
    # 保存数据到 csv
    def save_item(self,data):
        if self.dataMode == 1 :
            data.to_csv('data.csv', index=False) # 测试 保存数据到 csv 文件
            print('保存成功')
        elif self.dataMode == 2:
            try:
                conn = sqlite3.connect(self.filename)
                table_name = 'users'  # 假设表名为 'people'
                data.to_sql(table_name, conn, if_exists='replace', index=False)
                conn.close()
                print('保存到数据库成功')
            except Exception as e:
                print(f"保存到数据库时发生错误: {e}")